package top.yudoge.dogeeditor.llm.util

import dev.langchain4j.data.embedding.Embedding
import dev.langchain4j.store.embedding.EmbeddingSearchRequest
import top.yudoge.dogeeditor.llm.embedding.EmbeddingStoreHolder

fun EmbeddingStoreHolder.DeltaEmbeddingStore.getRelatedTableNames(
    searchEmbedding: Embedding,
    topK: Int,
    embeddingStoreHolder: EmbeddingStoreHolder
): List<String> {
    val request = EmbeddingSearchRequest(searchEmbedding, topK, 0.4, null)
    val results = this.search(request)
    return results.matches().map { embeddingStoreHolder.getTableName(it.embeddingId()) }
}